import paramiko
import threading
from os import path
from models import ConnectInfo

class SSHConnect:
    def __init__(self, connectInfo: ConnectInfo) -> None:
        self.connInfo = connectInfo
        self.ssh_client = paramiko.SSHClient()
        self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.status = False
    
    def connect(self):
        try:
            self.ssh_client.connect(
                hostname=self.connInfo.hostname,
                port=int(self.connInfo.port),
                username=self.connInfo.username,
                password=self.connInfo.pwd,
                timeout=10)
            self.status = True
            print(f"✓ Successfully connected to {self.connInfo.name}")
        except Exception as e:
            self.status = False
            print(f"✗ Failed to connect to {self.connInfo.name}: {e}")
    
    def command(self, command: str) -> bool:
        if not self.status:
            print(f"Not connected to {self.connInfo.name}, attempting to connect...")
            self.connect()
            if not self.status:
                return False
        
        try:
            stdin, stdout, stderr = self.ssh_client.exec_command(command=command, timeout=30)
            exit_status = stdout.channel.recv_exit_status()
            
            stdout_output = stdout.read().decode('utf-8').strip()
            stderr_output = stderr.read().decode('utf-8').strip()
            
            if stdout_output:
                print(f"[{self.connInfo.name} - STDOUT]: {stdout_output}")
            if stderr_output:
                print(f"[{self.connInfo.name} - STDERR]: {stderr_output}")
            
            return exit_status == 0
            
        except Exception as e:
            print(f"Error executing command on {self.connInfo.name}: {e}")
            return False
        
    def fileupload(self, inpath: str, distpath: str):
        print(inpath)
        print(distpath)
        if not path.exists(inpath):
            print("The file is not exist!")
            return False
        
        try:
            sftp = self.ssh_client.open_sftp()
            sftp.put(inpath, distpath)
            sftp.close()
            print(f"✓ Successfully uploaded {inpath} to {distpath} on {self.connInfo.name}")
            return True
        except Exception as e:
            print(f"✗ Failed to upload file to {self.connInfo.name}: {e}")
            return False

    def close(self):
        """Close the SSH connection"""
        if self.ssh_client:
            self.ssh_client.close()
            self.status = False

    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()